
#include <aftermath/probability/dist_exponential.hpp>
#include <aftermath/probability/dist_normal.hpp>
#include <aftermath/random/sampler_bernoulli.hpp>
#include <aftermath/random/sampler_binomial_lookup.hpp>
#include <aftermath/random/sampler_exponential.hpp>
#include <aftermath/random/sampler_lognormal.hpp>
#include <aftermath/random/sampler_pareto.hpp>
#include <aftermath/random/sampler_normal.hpp>
#include <aftermath/random/ziggurat_exponential_1024_uint32.hpp>
#include <aftermath/random/ziggurat_normal_1024_uint32.hpp>

#include <chrono>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <random>
#include <string>
#include <type_traits>

template <typename t_distribution_type>
struct sampler_switch
{
    typedef t_distribution_type type;
};

template <>
struct sampler_switch<aftermath::probability::dist_normal>
{
    typedef aftermath::random::sampler_normal_from_engine<1024, std::default_random_engine>::type type;
};

template <>
struct sampler_switch<aftermath::probability::dist_exponential>
{
    typedef aftermath::random::sampler_exponential_from_engine<1024, std::default_random_engine>::type type;
};

template <>
struct sampler_switch<aftermath::probability::dist_lognormal>
{
    typedef aftermath::random::sampler_lognormal_from_engine<1024, std::default_random_engine>::type type;
};

template <typename t_distribution_type>
void test_tail(std::uint_fast64_t n, double tail, t_distribution_type&& dist)
{
    typedef typename sampler_switch<t_distribution_type>::type sampler_type;
    typedef std::default_random_engine engine_type;
    std::cout << "Testing tail probability:  ";

    sampler_type smp = sampler_type(dist);
    std::uint_fast32_t seed = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    engine_type engine(seed);

    std::uint_fast64_t count_tail = 0UL;
    for (std::uint_fast64_t i = 0UL; i < n; i++)
    {
        if (smp(engine) > tail) count_tail++;
    }
    double p_tail = static_cast<double>(count_tail) / static_cast<double>(n);
    std::cout << "P(X > " << tail << ") = " << p_tail << std::endl;
}

template <typename t_distribution_type>
void test_cusum(std::uint_fast64_t n, double threshold, t_distribution_type&& dist)
{
    typedef typename sampler_switch<t_distribution_type>::type sampler_type;
    typedef std::default_random_engine engine_type;
    std::cout << "Testing CUSUM average run length:  ";

    sampler_type smp = sampler_type(dist);
    std::uint_fast32_t seed = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    engine_type engine(seed);

    double drift = dist.mean() + dist.stddev() / 2;

    std::cout << std::endl;
    double sum = 0.0;
    for (std::uint_fast64_t i = 0UL; i < n; i++)
    {
        std::uint_fast64_t time = 0;
        double cusum_ref = 0.0;
        while (cusum_ref < threshold)
        {
            double z = smp(engine) - drift;
            cusum_ref = (cusum_ref < 0.0 ? 0.0 : cusum_ref) + z;

            time++;
        }
        sum += time;
    }

    std::cout << "ARL = " << sum / n << std::endl;
}

template <std::size_t n_min, std::size_t n_max, typename t_sampler_type>
void test_binomial(std::size_t m, double p)
{
    typedef std::chrono::high_resolution_clock clock_type;
    auto tic = clock_type::now();

    std::uniform_int_distribution<std::int32_t> uniform_n = std::uniform_int_distribution<std::int32_t>(n_min, n_max);
    //std::uniform_real_distribution<double> uniform_p = std::uniform_real_distribution<double>(0.0, 1.0);
    std::uint_fast32_t seed = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    std::default_random_engine engine = std::default_random_engine(seed);

    std::size_t sum = 0UL;
    for (std::size_t i = 0UL; i < m; i++)
    {
        auto n = uniform_n(engine);
        aftermath::probability::dist_binomial distribution(n, p);
        t_sampler_type sampler(distribution);

        sum += sampler(engine);
    }

    double mean = static_cast<double>(sum) / m;
    std::cout << "Binomial(Unif[" << n_min << ", " << n_max << "], " << p << "):" << std::endl;
    std::cout << "----Mean = " << mean << std::endl;

    auto toc = clock_type::now();
    auto dt = toc.time_since_epoch() - tic.time_since_epoch();
    std::cout << "Elapsed time " << static_cast<double>(dt.count() * clock_type::period::num) / clock_type::period::den << " s." << std::endl;
}

template <std::size_t n_min, std::size_t n_max>
void test_binomial_table(std::size_t m, double p)
{
    typedef std::chrono::high_resolution_clock clock_type;
    auto tic = clock_type::now();

    std::uniform_int_distribution<std::int32_t> uniform_n = std::uniform_int_distribution<std::int32_t>(n_min, n_max);
    //std::uniform_real_distribution<double> uniform_p = std::uniform_real_distribution<double>(0.0, 1.0);
    std::uint_fast32_t seed = static_cast<std::uint_fast32_t>(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    std::default_random_engine engine = std::default_random_engine(seed);

    auto from = aftermath::probability::dist_binomial(n_min, p);
    auto to = aftermath::probability::dist_binomial(n_max, p);
    typename aftermath::random::sampler_binomial_lookup_from_engine<std::default_random_engine>::type binomial_matrix(from, to);
    std::cout << binomial_matrix.height() << "-by-" << binomial_matrix.width() << " (" << (binomial_matrix.size_in_bytes() / (1024.0 * 1024.0)) << " Mb)..." << std::endl;

    std::size_t sum = 0UL;
    for (std::size_t i = 0UL; i < m; i++)
    {
        auto n = uniform_n(engine);
        //binomial_matrix.at(n);

        sum += binomial_matrix(n, engine);
    }

    double mean = static_cast<double>(sum) / m;
    std::cout << "Binomial(Unif[" << n_min << ", " << n_max << "], " << p << "):" << std::endl;
    std::cout << "----Mean = " << mean << std::endl;

    auto toc = clock_type::now();
    auto dt = toc.time_since_epoch() - tic.time_since_epoch();
    std::cout << "Elapsed time " << static_cast<double>(dt.count() * clock_type::period::num) / clock_type::period::den << " s." << std::endl;
}

template <typename F>
void benchmark(F func)
{
    typedef std::chrono::high_resolution_clock clock_type;
    auto tic = clock_type::now();
    func();
    auto toc = clock_type::now();
    auto dt = toc.time_since_epoch() - tic.time_since_epoch();

    std::cout << "Elapsed time " << static_cast<double>(dt.count() * clock_type::period::num) / clock_type::period::den << " s." << std::endl;
}

int main() // (int argc, char* argv[], char* envp[])
{
    typedef aftermath::probability::dist_normal    aftermath_normal;
    typedef aftermath::probability::dist_lognormal aftermath_lognormal;

    double mean = 0.0;
    double variance = 1.0;
    double sigma = std::sqrt(variance);
    double tail = 3.0;
    double threshold = std::log(159.36);
    
    const std::uint_fast64_t n1 = 10000000UL; // Tail.
    const std::uint_fast64_t n2 = 100000UL; // CUSUM.

    std::cout << "c++11 builtin gaussian:" << std::endl;
    benchmark([n1, tail, mean, sigma]() { test_tail(n1, tail, std::normal_distribution<double>(mean, sigma)); });
    benchmark([n2, threshold, mean, sigma]() { test_cusum(n2, threshold, std::normal_distribution<double>(mean, sigma)); });

    std::cout << "aftermath gaussian:" << std::endl;
    benchmark([n1, tail, mean, sigma]() { test_tail(n1, tail, aftermath_normal(mean, sigma)); });
    benchmark([n2, threshold, mean, sigma]() { test_cusum(n2, threshold, aftermath_normal(mean, sigma)); });

    std::cout << "c++11 builtin lognormal:" << std::endl;
    benchmark([n1, tail, mean, sigma]() { test_tail(n1, tail, std::lognormal_distribution<double>(mean, sigma)); });

    std::cout << "aftermath lognormal:" << std::endl;
    benchmark([n1, tail, mean, sigma]() { test_tail(n1, tail, aftermath_lognormal(mean, sigma)); });

    // ~~ Binomial ~~
    const std::size_t n_trials_min = 1;
    const std::size_t n_trials_max = 299;
    double p = 0.8;

    std::cout << "Bernoulli:" << std::endl;
    test_binomial<n_trials_min, n_trials_max, typename aftermath::random::sampler_bernoulli_from_engine<std::default_random_engine>::type>(n1, p);

    std::cout << "Consolidated table:" << std::endl;
    test_binomial_table<n_trials_min, n_trials_max>(n1, p);

    //std::cout << "Press any key to continue . . ." << std::endl;
    return 0;
}
